# -*- coding: utf-8 -*-
import os
import torch
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling,
)

# 基础模型配置
model_name = "/opt/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" # 预训练模型名称或本地模型路径文件夹
save_dir = "./private_deepseek_qwen" # 微调后模型保存路径
data_path = "./private_data.jsonl" # JSONL格式的私有训练数据

# GPU检测
use_gpu = torch.cuda.is_available()
device = torch.device("cuda" if use_gpu else "cpu")
print(f"检测到设备: {device}, {'使用GPU' if use_gpu else '使用CPU'}")

# 加载预训练模型和分词器
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16 if use_gpu else torch.float32,  # GPU使用半精度浮点数
    device_map="auto" if use_gpu else "cpu",  # 自动处理多GPU部署
    trust_remote_code=True
)

# 数据预处理函数 - 将文本转换为模型可接受的输入格式
def preprocess_function(examples):
    # 对文本进行分词并填充到固定长度
    inputs = tokenizer(examples["text"], truncation=True, max_length=512, padding="max_length")
    inputs["labels"] = inputs["input_ids"].copy()  # 语言模型的标签与输入相同
    return inputs

# 加载和准备训练数据集
dataset = load_dataset("json", data_files=data_path)  # 加载JSONL格式数据
tokenized_dataset = dataset.map(preprocess_function, batched=True)  # 批量处理数据

# 训练参数配置
training_args = TrainingArguments(
    output_dir=save_dir,  # 模型保存目录
    learning_rate=2e-5,  # 学习率 - 控制参数更新的步长
    per_device_train_batch_size=4,  # 每个设备的训练批次大小
    per_device_eval_batch_size=4,  # 每个设备的评估批次大小
    num_train_epochs=3,  # 训练轮数 - 整个数据集训练的次数
    weight_decay=0.01,  # 权重衰减 - 防止过拟合的正则化方法
    save_strategy="steps",  # 保存策略 - 按步数保存模型
    save_steps=1000,  # 每1000步保存一次模型
    push_to_hub=False,  # 是否推送到Hugging Face模型库
    fp16=use_gpu,  # 是否使用混合精度训练
    logging_dir="./logs",  # 日志保存目录
    logging_steps=10,  # 每10步记录一次日志
    report_to="all",  # 报告训练指标到所有支持的平台
    dataloader_pin_memory=use_gpu,  # 仅在GPU时启用内存锁定以加速数据传输
)

# 数据收集器 - 处理批次数据的组装和填充
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, 
    mlm=False,  # 不使用掩码语言模型任务
    pad_to_multiple_of=8  # 优化内存对齐，提高GPU利用率
)

# 初始化训练器
trainer = Trainer(
    model=model,  # 待训练的模型
    args=training_args,  # 训练参数
    train_dataset=tokenized_dataset["train"],  # 训练数据集
    eval_dataset=None,  # 不使用验证集
    data_collator=data_collator,  # 数据收集器
)

print(f"执行训练...")
# 执行训练
trainer.train()
print(f"训练结束")
# 保存模型和分词器
trainer.save_model(save_dir)
tokenizer.save_pretrained(save_dir)    
print(f"保存完成")